#import os
import sys
import numpy as np
import matplotlib.pyplot as plt

"""
from optparse import OptionParser
parser = OptionParser("usage:  %prog [options] arg1 ")
parser.add_option("-k", "--kmcolumn", dest="kmcolumn")
parser.add_option("-t", "--transient", dest="transient")
"""


### Particle class
class Particle2D:
  def __init__(self):
    self.x_new = 0.0
    self.y_new = 0.0
    self.vx_new = 0.0
    self.vy_new = 0.0
    self.fx_new = 0.0
    self.fy_new = 0.0
    self.save_old()

  def save_old(self):
    self.x_old  = self.x_new
    self.y_old  = self.y_new
    self.vx_old = self.vx_new
    self.vy_old = self.vy_new
    self.fx_old = self.fx_new
    self.fy_old = self.fy_new
    
    
##################################################################
### Gas class 
# TODO: define boundary conditions
class Gas2D:
  #--- Box init dimensions, N
  def __init__(self, xlen=100, ylen=10, N=100, KT=0.1, delta_r=1.0, iniFile=""):
    self.xlen = xlen
    self.ylen = ylen
    if iniFile=="":
      self.N = N
      self.particle = [ Particle2D() for i in range(self.N)]  
      self.KT = KT
    else:
      data = np.loadtxt(iniFile)
      self.N = len(data)
      self.particle = [ Particle2D() for i in range(self.N)]  
      self.KT = KT
      
    
    self.X = np.zeros(self.N)
    self.Y = np.zeros(self.N)
    self.VX = np.zeros(self.N)
    self.VY = np.zeros(self.N)
    #self.R2  = np.zeros((self.N,self.N))
    self.Ri2 = np.zeros((self.N,self.N))
    self.Ri6 = np.zeros((self.N,self.N))
    self.Ri12 = np.zeros((self.N,self.N))
    self.Fx_new = np.zeros((self.N,self.N))
    self.Fy_new = np.zeros((self.N,self.N))
    self.Fx_old = np.zeros((self.N,self.N))
    self.Fy_old = np.zeros((self.N,self.N))
    self.Epot = 0.0
    self.Ekin = 0.0
    self.Etot = 0.0
    self.delta_r = delta_r
    
    
    if iniFile=="": 
      self.initPos()
      self.initVel(KT)
      self.speed  = np.zeros(self.N)
    else:
        for i in range(self.N):
          self.particle[i].x_new=data[i][0]
          self.particle[i].y_new=data[i][1]
          self.particle[i].vx_new=data[i][2]
          self.particle[i].vy_new=data[i][3]
    
  def save_old(self):
    for pi in self.particle:
      pi.save_old()
  #--- Initial Positions
  def initPos(self):
    # Uniform random positions
    #for pi in self.particle:
    #  pi.x_new = np.random.uniform(0, self.xlen)
    #  pi.y_new = np.random.uniform(0, self.ylen)

    self.particle[0].x_new = self.delta_r
    self.particle[0].y_new = self.delta_r
    
    for i in range(1,self.N):
      self.particle[i].x_new = self.particle[i-1].x_new + self.delta_r
      self.particle[i].y_new = self.particle[i-1].y_new
      if (self.particle[i].x_new + self.delta_r) > self.xlen :
        self.particle[i].x_new = self.delta_r
        self.particle[i].y_new = self.particle[i].y_new + self.delta_r
        if self.particle[i].y_new > self.ylen:
          print 'rescale sigma or delta_r', i
          sys.exit()
          
  def initPos02(self):
    self.particle[0].x_new = self.delta_r
    self.particle[0].y_new = self.delta_r
    for i in range(1,self.N):
      self.particle[i].x_new = self.particle[i-1].x_new + self.delta_r
      self.particle[i].y_new = self.particle[i-1].y_new
    
  #--- Initial Velocities
  def initVel(self,KT):
    # Maxwell-Boltzmann velocities distribution
    np.random.seed(10)
    for pi in self.particle:
      pi.vx_new = np.random.normal(0.0, np.sqrt(KT))
      pi.vy_new = np.random.normal(0.0, np.sqrt(KT))
  
  #--- Calculate R^-2 matrix
  def calculateRi2(self):
    aux = 0.0
    for i in range(self.N):
      for j in range(i,self.N):
        if i == j :
          self.Ri2[i][j] = 0.0
        else :
          aux = (self.particle[i].x_new - self.particle[j].x_new)**2 + (self.particle[i].y_new - self.particle[j].y_new)**2
          #print i,j, aux
          self.Ri2[i][j] = 1.0/aux
          self.Ri2[j][i] = self.Ri2[i][j]

  ####### Lennard-Jones #######
  #--- Calculate the force with R2 matrix
  def calculateForces(self, eps, sigma, exclude=[]):
    self.calculateRi2()
    self.Ri6  = self.Ri2*self.Ri2*self.Ri2
    self.Ri12 = self.Ri6*self.Ri6
    self.Epot = 0.0
    sigma6  = sigma**6
    sigma12 = sigma6**2
    aux = 0.0
    
    rc=2.5
    rci2  = 1/(rc*rc)
    rci6  = rci2*rci2
    rci12 = rci6*rci6
    C2 = 12.0*eps*( 2.0*sigma12*rci12  - sigma6*rci6 )
    UC =  4.0*eps*(-7.0*sigma12*rci12 + 4.0*sigma6*rci6)
    
    for i in range(self.N):
      for j in range(i+1,self.N):
        if (i in exclude):
          # TODO: keep the position fixed!!
          print 'exclude'
        else:
          # calculate the force
          if self.Ri2[i][j] >= rci2: # R^2 <= rc^2 so Ri2 >= rci2 (R, rc inverse)
            # Potential LJ energia new
            self.Epot = self.Epot + 4*eps*(sigma12*self.Ri12[i][j] - sigma6*self.Ri6[i][j]) + C2*self.Ri2[i][j]*rci2 + UC
            # Force LJ new
            aux = 24*eps*(2*sigma12*self.Ri12[i][j] - sigma6*self.Ri6[i][j])*self.Ri2[i][j] - 2*C2*rci2
            deltaX = (self.particle[i].x_new - self.particle[j].x_new)
            deltaY = (self.particle[i].y_new - self.particle[j].y_new)

            self.Fx_new[i][j] = aux*deltaX 
            self.Fx_new[j][i] = -self.Fx_new[i][j]
            
            self.Fy_new[i][j] = aux*deltaY
            self.Fy_new[j][i] = -self.Fy_new[i][j]
          else:
            self.Fx_new[i][j] = 0.0
            self.Fx_new[j][i] = 0.0
            
            self.Fy_new[i][j] = 0.0
            self.Fy_new[j][i] = 0.0
        

    for i in range(self.N):
      self.particle[i].fx_new = self.Fx_new[i].sum()
      self.particle[i].fy_new = self.Fy_new[i].sum()
      
  
  ######## Heat bath ########
  #--- Calculate the force with R2 matrix for a 1-D chain
  def calculateForces_HeatBath(self, eps, sigma,T1,T2):
    self.calculateRi2()
    self.Ri6  = self.Ri2*self.Ri2*self.Ri2
    self.Ri12 = self.Ri6*self.Ri6
    self.Epot = 0.0
    sigma6  = sigma**6
    sigma12 = sigma6**2
    aux = 0.0
    
    rc=2.5
    rci2  = 1/(rc*rc)
    rci6  = rci2*rci2
    rci12 = rci6*rci6
    C2 = 12.0*eps*( 2.0*sigma12*rci12  - sigma6*rci6 )
    UC =  4.0*eps*(-7.0*sigma12*rci12 + 4.0*sigma6*rci6)
    
    for i in range(self.N):
      for j in range(i+1,self.N):
        # calculate the force
        if self.Ri2[i][j] >= rci2: # R^2 <= rc^2 so Ri2 >= rci2 (R, rc inverse)
          # Potential LJ energia new
          self.Epot = self.Epot + 4*eps*(sigma12*self.Ri12[i][j] - sigma6*self.Ri6[i][j]) + C2*self.Ri2[i][j]*rci2 + UC
          # Force LJ new
          aux = 24*eps*(2*sigma12*self.Ri12[i][j] - sigma6*self.Ri6[i][j])*self.Ri2[i][j] - 2*C2*rci2
          deltaX = (self.particle[i].x_new - self.particle[j].x_new)
          deltaY = (self.particle[i].y_new - self.particle[j].y_new)

          self.Fx_new[i][j] = aux*deltaX 
          self.Fx_new[j][i] = -self.Fx_new[i][j]
          
          self.Fy_new[i][j] = aux*deltaY
          self.Fy_new[j][i] = -self.Fy_new[i][j]
        else:
          self.Fx_new[i][j] = 0.0
          self.Fx_new[j][i] = 0.0
          
          self.Fy_new[i][j] = 0.0
          self.Fy_new[j][i] = 0.0
        
    for i in range(self.N):
      self.particle[i].fx_new = self.Fx_new[i].sum()
      self.particle[i].fy_new = self.Fy_new[i].sum()
  
    
    ### Wall interactions like LJ and reservoir (thermic bath)
    gamma = 1.0
    #self.particle[0]
    R_W0    = self.particle[0].x_new - 0.0
    Ri2_W0  = 1/(R_W0**2)
    if Ri2_W0 >= rci2:
      Ri6_W0  = Ri2_W0*Ri2_W0*Ri2_W0
      Ri12_W0 = Ri6_W0*Ri6_W0
      
      aux_W0 = 24*eps*(2*sigma12*Ri12_W0 - sigma6*Ri6_W0)*Ri2_W0 - 2*C2*rci2
      self.particle[0].fx_new += aux_W0*R_W0
      # the reservoir
      self.particle[0].fx_new += -gamma*self.particle[0].vx_new
      self.particle[0].fx_new += np.sqrt(2*T1)*np.random.normal(0,1)
    
    #self.particle[self.N-1]
    R_LW    = self.xlen - self.particle[self.N-1].x_new # FIXME: Possible ERROR
    Ri2_LW  = 1/(R_LW**2)
    if Ri2_W0 >= rci2:
      Ri6_LW  = Ri2_LW*Ri2_LW*Ri2_LW
      Ri12_LW = Ri6_LW*Ri6_LW
      
      aux_LW = 24*eps*(2*sigma12*Ri12_LW - sigma6*Ri6_LW)*Ri2_LW - 2*C2*rci2
      self.particle[self.N-1].fx_new += aux_LW*R_LW
      # the reservoir
      self.particle[self.N-1].fx_new += -gamma*self.particle[self.N-1].vx_new
      self.particle[self.N-1].fx_new += np.sqrt(2*T1)*np.random.normal(0,1)
      
  ####### Harmonic chain #######
  def calculateForceHarmonic(self,len0):
    # fixed chain extremes 
    for i in range(1,self.N-1):  
      # Calculate the force:
      self.particle[i].fx_new = (self.particle[i+1].x_new - self.particle[i].x_new - len0) - (self.particle[i].x_new - self.particle[i-1].x_new - len0)
      self.particle[i].fy_new = 0.0
    
  #--- Validate Particle in the box
  # FIXME: overlap is possible!! How to avoid that?
  def inBox(self):
    for pi in self.particle:
      delta_x =  self.xlen - pi.x_new
      delta_y =  self.ylen - pi.y_new
      if pi.x_new < 0:
        pi.x_new  = -pi.x_new
        pi.vx_old = -pi.vx_old
      elif delta_x < 0 :
        # A Wall collision occurs in x axis
        pi.x_new  = self.xlen + delta_x
        pi.vx_old = -pi.vx_old # Because the vx_new will use vx_old
        
      if pi.y_new < 0:
        pi.y_new  = -pi.y_new
        pi.vy_old = -pi.vy_old
      elif delta_y < 0:
        # A Wall collision occurs in x axis
        #print "--- Y up collision ", pi.y_new, pi.vy_new, pi.vx_new, delta_y
        pi.y_new  = self.ylen + delta_y
        pi.vy_old = - pi.vy_old # Because the vy_new will use vy_old
  
  #--- Step position
  def stepPosition(self, dt):
    for pi in self.particle:
      pi.x_new = pi.x_old + dt*pi.vx_old + 0.5*dt*dt*(pi.fx_old)
      pi.y_new = pi.y_old + dt*pi.vy_old + 0.5*dt*dt*(pi.fy_old)
      
  #--- Step velocity
  def stepVelocity(self, dt):
    for pi in self.particle:
      #print 'force:', pi.fx_new, pi.fx_old, self.Epot
      pi.vx_new = pi.vx_old + 0.5*dt*(pi.fx_new + pi.fx_old)
      pi.vy_new = pi.vy_old + 0.5*dt*(pi.fy_new + pi.fy_old)
      #print 'velocity:',pi.vx_new, pi.vx_old
      
  #--- Generate two arrays for the positions.      
  def arrayPos(self):
    self.X =np.array([self.particle[i].x_new for i in range(self.N)])
    self.Y =np.array([self.particle[i].y_new for i in range(self.N)])

  #--- Generate two arrays for the velocities.
  def arrayVel(self):
    self.VX =np.array([self.particle[i].vx_new for i in range(self.N)])
    self.VY =np.array([self.particle[i].vy_new for i in range(self.N)])
  
  def calculateEkin(self):
    self.arrayVel()
    self.Ekin = 0.5*((self.VX**2 + self.VY**2).sum())
  #--- Generate speed array  
  def arraySpeed(self):
    self.speed = np.sqrt(self.VX**2+self.VY**2)
    
  #------------- Save and load data Functions --------
  def savetxt(self,outdatafile):
    self.arrayPos()
    self.arrayVel()
    np.savetxt(outdatafile,zip(self.X, self.Y, self.VX, self.VY))
    print outdatafile
    
  def loadtxt(self,inputdatafile):
    data = np.loadtxt(inputdatafile)
    if len(data) != self.N:
      print "not the same particles"
      sys.exit()
    else:
      for i in range(self.N):
        self.particle[i].x_new=data[i][0]
        self.particle[i].y_new=data[i][1]
        self.particle[i].vx_new=data[i][2]
        self.particle[i].vy_new=data[i][3]
      print inputdatafile
        
  #------------- Plot Functions --------
  def plot(self):
    my_dpi=96
    plt.rc('font', size=25)
    plt.rc('text', usetex=True)
    fig = plt.figure(figsize=(800/my_dpi, 800/my_dpi), dpi=my_dpi)
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlim((0,self.xlen))
    ax.set_ylim((0,self.ylen))
    for i in range(self.N):
      circle=plt.Circle((self.particle[i].x_new,self.particle[i].y_new),0.5,color='r',alpha=0.5)
      ax.add_patch(circle)
    self.arrayPos()
    self.arrayVel()
    ax.set_aspect('equal')
    ax.quiver(self.X, self.Y, self.VX, self.VY, color='b', pivot='mid', alpha =0.5, units='inches')
    plt.scatter(self.X, self.Y)
    plt.xlabel("$x$")
    plt.ylabel("$y$")
    plt.grid()
    
    plt.savefig("figura.pdf",dpi=100)
    plt.show()
    
  def hist(self):
    plt.figure()
    self.arrayVel()
    self.arraySpeed()
    plt.hist(self.speed)
    plt.show()
    
    
##################################################################
### GasTime class  
class GasTime:
  def __init__(self, gas, timeDim=1):
    self.timeDim = timeDim
    self.time  = np.zeros(self.timeDim)
    self.Ekint = np.zeros(self.timeDim)
    self.Epott = np.zeros(self.timeDim)
    self.Etott = np.zeros(self.timeDim)
    self.fluxt = np.zeros(self.timeDim)
    # could be pressure
    self.time_i = 0
    
    #Micro states in time
    self.xt  = np.zeros([gas.N, timeDim])
    self.yt  = np.zeros([gas.N, timeDim])
    self.vxt = np.zeros([gas.N, timeDim])
    self.vyt = np.zeros([gas.N, timeDim])
    self.fxt = np.zeros([gas.N, timeDim])
    self.fyt = np.zeros([gas.N, timeDim])
    
  def saveState(self, gas, step, dt):
    self.time[self.time_i]=dt*step
    self.Ekint[self.time_i] = gas.Ekin
    self.Epott[self.time_i] = gas.Epot
    self.Etott[self.time_i] = gas.Etot
    self.time_i = self.time_i + 1
    print self.time_i,
  
  def saveMicroState(self, gas, step, dt):
    for j in range(gas.N):
      self.xt[j][self.time_i]  = gas.particle[j].x_new
      self.yt[j][self.time_i]  = gas.particle[j].y_new
      self.vxt[j][self.time_i] = gas.particle[j].vx_new
      self.vyt[j][self.time_i] = gas.particle[j].vy_new
      self.fxt[j][self.time_i] = gas.particle[j].fx_new
      self.fyt[j][self.time_i] = gas.particle[j].fy_new

  def savetxt(self,filename):
    #TODO: save in a file the gasTime object
    print "fuck u"+filename
